# -*- coding: utf8 -*-

import cv2
import torch
from ultralytics import YOLO
from od_config import *

model_all_target = YOLO(weights_target)
model_plugging = YOLO(weights_plugging)


def detect_tail_fiber(image,model,minScore):
    results = model(image, imgsz=640,conf=minScore, iou=0.5,half=True)
    boxes = results[0].boxes
    names = results[0].names
    output = []
    for box in boxes:
        xy_box = box.xyxy.tolist()[0]
        box_cls = box.cls.tolist()[0]
        box_cof = round(box.conf.tolist()[0], 2)
        box_out = [int(xy_box[0]), int(xy_box[1]), int(xy_box[2]), int(xy_box[3]), box_cof, names[box_cls]]
        output.append(box_out)
    return output


if __name__ == '__main__':
    img_path = 'test_images/car_new.png'
    image = cv2.imread(img_path)
    print('image.shape:', image.shape)
    with torch.no_grad():
        x = detect_tail_fiber(image,model=model_plugging,minScore=0.5)
        print(x)
